package com.example.servicegateway.loadbalancer;

import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.ListUtils;
import org.springframework.boot.convert.DurationStyle;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.discovery.DiscoveryClient;
import org.springframework.cloud.client.discovery.ReactiveDiscoveryClient;
import org.springframework.cloud.client.loadbalancer.Request;
import org.springframework.cloud.client.loadbalancer.RequestDataContext;
import org.springframework.cloud.loadbalancer.core.ServiceInstanceListSupplier;
import org.springframework.core.env.Environment;
import reactor.core.publisher.Flux;
import reactor.core.scheduler.Schedulers;

import java.time.Duration;
import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

import static org.springframework.cloud.loadbalancer.support.LoadBalancerClientFactory.PROPERTY_NAME;

/**
 * 定义选取灰度服务列表Supplier
 * @author K
 */
@Slf4j
public class GrayServiceInstanceListSupplier implements ServiceInstanceListSupplier {

    /**
     * Property that establishes the timeout for calls to service discovery.
     */
    public static final String SERVICE_DISCOVERY_TIMEOUT = "spring.cloud.loadbalancer.service-discovery.timeout";

    private Duration timeout = Duration.ofSeconds(30);

    private final String serviceId;

    private final Flux<List<ServiceInstance>> serviceInstances;

    /**
     * 使用@LoadBalancerClient指定服务使用负载均衡器
     * @param delegate 服务发现客户端
     * @param environment 环境变量对象
     */
    public GrayServiceInstanceListSupplier(DiscoveryClient delegate, Environment environment) {
        this.serviceId = environment.getProperty(PROPERTY_NAME);
        resolveTimeout(environment);
        this.serviceInstances = Flux.defer(() -> Flux.just(delegate.getInstances(serviceId)))
                .subscribeOn(Schedulers.boundedElastic()).timeout(timeout, Flux.defer(() -> {
                    logTimeout();
                    return Flux.just(new ArrayList<>());
                })).onErrorResume(error -> {
                    logException(error);
                    return Flux.just(new ArrayList<>());
                });
    }


    public GrayServiceInstanceListSupplier(ReactiveDiscoveryClient delegate, Environment environment) {
        this.serviceId = environment.getProperty(PROPERTY_NAME);
        resolveTimeout(environment);
        this.serviceInstances = Flux
                .defer(() -> delegate.getInstances(serviceId).collectList().flux().timeout(timeout, Flux.defer(() -> {
                    logTimeout();
                    return Flux.just(new ArrayList<>());
                })).onErrorResume(error -> {
                    logException(error);
                    return Flux.just(new ArrayList<>());
                }));
    }

    @Override
    public String getServiceId() {
        return serviceId;
    }

    /**
     * 主要重写方法
     */
    @Override
    public Flux<List<ServiceInstance>> get(Request request) {
        final String gray = ((RequestDataContext) request.getContext()).getClientRequest().getHeaders().getFirst("gray");
        if (gray != null){
            //serviceInstances Flux在subscribe时才会执行map代码，所以我们只能将容错代码写在map内
            return serviceInstances.map(list ->{
                List<ServiceInstance> tmpList = list.stream()
                        .filter(instance -> Boolean.parseBoolean(instance.getMetadata().get("gray")))
                        .collect(Collectors.toList());

                if (Boolean.parseBoolean(gray)) {
                    return tmpList.size() > 0 ? tmpList : list;
                }
                //容错，如header无匹配的则返回正常列表轮询
                return ListUtils.removeAll(list, tmpList);
            });
        }
        return serviceInstances;
    }

    @Override
    public Flux<List<ServiceInstance>> get() {
        return serviceInstances;
    }

    private void resolveTimeout(Environment environment) {
        String providedTimeout = environment.getProperty(SERVICE_DISCOVERY_TIMEOUT);
        if (providedTimeout != null) {
            timeout = DurationStyle.detectAndParse(providedTimeout);
        }
    }

    private void logTimeout() {
        log.info(String.format("Timeout occurred while retrieving instances for service %s."
                + "The instances could not be retrieved during %s", serviceId, timeout));
    }

    private void logException(Throwable error) {
        log.warn(String.format("Exception occurred while retrieving instances for service %s", serviceId), error);
    }
}
